#  -*- coding: utf-8 -*-

from strategy.stock_pool.base_stock_pool import BaseStockPool
from factor.factor_module import FactorModule
from data.data_module import DataModule
from util.stock_util import judge_code_trading_date,get_all_codes_date,get_diff_dates,calc_negative_diff_dates,multi_computer,get_all_codes_trading_date,get_code_name
from pandas import DataFrame,Series
import pandas as pd
import matplotlib.pyplot as plt
import numpy as np
import pylab as pl
import scipy.signal as signal
import time
from sklearn import linear_model
import pickle
import os
from util.database import DB_CONN
from pymongo import UpdateOne,ASCENDING
from pathlib import Path
import talib as ta
from datetime import datetime
from util.database import base_code_path

"""
价值股跟踪：
1，读取价值跟踪股的数据；
2，补充完成后，存入数据库；
3，在visual模块去展示部分数据
4，定期去更新之前一些字段
"""

class ValueStockPool(BaseStockPool):
    dm = DataModule()
    fm = FactorModule()
    cur_date = datetime.now().strftime('%Y-%m-%d')
    cur_date = judge_code_trading_date(date=cur_date)

    def read_basic_info(self):
        self.file = Path(f"{base_code_path}ManualStockPool.xlsx")

        data_df = pd.read_excel(self.file, sheet_name='ValueTracking', sheencoding="gb2312", dtype=object)
        #print(data_df.head())

        change_col = {
            "代码":'code',
            "名称":'name',
            "合理价格下限":'low_price',
            "合理价格上限":'high_price',
            "更新时间":"update_date",
            "来源":'source',
            "行业":'industry',
            "类别":'category',
            "分析":"analyze"
        }
        data_df.rename(columns=change_col,inplace=True)
        data_df['update_date'] = data_df['update_date'].astype(str)
        data_df['code'] = data_df['code'].astype(str)

        #print(data_df.head())

        return data_df

    def get_cur_price(self,dm,code,date):
        date_df = dm.get_k_data(code,autype='qfq',begin_date=date,end_date=date)
        close = date_df.loc[0]['close']
        return close


    def calc_bias_cur_price(self,cur_price,high_price):
        if high_price != 0:
            bias_ratio = round(100 * (cur_price - high_price) / high_price, 2)
        else:
            bias_ratio = 'N/A'
        return bias_ratio

    def fill_value(self,data_df):

        data_df['cur_price'] = data_df.apply(lambda row: self.get_cur_price(self.dm,row['code'],self.cur_date), axis=1)
        data_df['bias_cur_price_ratio'] = data_df.apply(lambda row:self.calc_bias_cur_price(row['cur_price'],row['high_price']),axis=1)

        return

    def write_value_data_to_db(self,data_df):

        update_requests = list()
        collection = DB_CONN['value_strategy_option_stocks']
        # 建立code+date的索引，提高save_data时写入数据的查询速度
        collection.create_index([('code', 1),('update_date', 1)])
        for index, row in data_df.iterrows():
            update_requests.append(
                UpdateOne(
                    {'index': index},
                    {'$set': dict(row)},
                    upsert=True)
            )
        # print(update_requests)
        # 批量写入，提高访问效率
        if len(update_requests) > 0:
            start_time = time.time()
            update_result = collection.bulk_write(update_requests, ordered=False)
            end_time = time.time()
            print('保存价值跟踪数据到数据集：%s，插入：%4d条, 更新：%4d条,耗时：%.3f 秒' %
                  (collection.name, update_result.upserted_count, update_result.modified_count,
                   (end_time - start_time)),
                  flush=True)

        return

    def read_value_data_from_db(self):
        collection = DB_CONN['value_strategy_option_stocks']

        data_cursor = collection.find(
            sort=[('update_date', ASCENDING)],
            projection={'_id': False})
        data_df = DataFrame([x for x in data_cursor])
        return data_df


    def get_option_stocks(self):
        flag = 1
        if flag == 1:
            #step1 从excel读取数据，并建立基础的data_df格式
            data_df = self.read_basic_info()

            #step2 填充data_df
            self.fill_value(data_df)

            #step3 写入数据库
            self.write_value_data_to_db(data_df)


        return

if __name__ == '__main__':
    pd.set_option('display.width', 130)
    pd.set_option('display.max_columns', 130)
    pd.set_option('display.max_colwidth', 130)

    vsp = ValueStockPool("Value",None,None,1)
    vsp.get_option_stocks()

